import curses
import random
import json
from typing import Dict, List, Tuple

from metagpt.actions import Action
from metagpt.roles import Role
from metagpt.schema import Message
from metagpt.logs import logger
from metagpt.llm import LLM

def clean_llm_response(response: str, expect_json: bool = False) -> str:
    """清理 LLM 返回的响应，移除 Markdown 格式"""
    # 移除所有代码块标记
    response = response.replace("```json", "").replace("```python", "").replace("```", "").strip()
    
    if expect_json:
        # 如果期望 JSON，尝试提取 JSON 部分
        try:
            start = response.find("{")
            end = response.rfind("}") + 1
            if start >= 0 and end > start:
                response = response[start:end]
        except:
            pass
    
    return response

class GameDesign(Action):
    """游戏设计动作：定义贪吃蛇游戏规则"""
    
    async def run(self, *args, **kwargs) -> Dict:
        """使用 LLM 生成游戏设计文档"""
        prompt = """请设计一个命令行贪吃蛇游戏的详细规范，包括：
1. 游戏名称和描述
2. 游戏规则（移动、得分、游戏结束条件等）
3. 主要功能特性
请以 JSON 格式返回，不要包含任何其他文本。"""
        
        llm = LLM()
        response = await llm.aask(prompt)
        response = clean_llm_response(response, expect_json=True)
        design_doc = json.loads(response)
        logger.info(f"Game design completed: {design_doc}")
        return design_doc

class CodeImplementation(Action):
    """代码实现动作：根据设计编写游戏代码"""
    
    async def run(self, design: Dict, *args, **kwargs) -> str:
        """使用 LLM 实现贪吃蛇游戏代码"""
        prompt = f"""根据以下游戏设计规范，实现一个完整的命令行贪吃蛇游戏：
{json.dumps(design, indent=2)}

要求：
1. 使用 Python 的 curses 库实现
2. 包含完整的游戏逻辑
3. 代码要清晰易读
4. 包含必要的错误处理

请直接返回完整的 Python 代码，不要包含任何其他文本。"""
        
        llm = LLM()
        code = await llm.aask(prompt)
        code = clean_llm_response(code)
        logger.info("Code implementation completed")
        return code

class GameTester(Action):
    """游戏测试动作：验证游戏功能"""
    
    async def run(self, code: str, *args, **kwargs) -> Dict:
        """使用 LLM 测试游戏代码"""
        prompt = f"""请分析以下贪吃蛇游戏代码，并验证其功能：
{code}

请检查：
1. 移动控制是否正常
2. 得分系统是否正确
3. 碰撞检测是否准确
4. 渲染是否正常

请以 JSON 格式返回测试结果，包含每个检查项的状态（PASS/FAIL），不要包含任何其他文本。
示例格式：
{{
    "movement": "PASS",
    "scoring": "PASS",
    "collision": "PASS",
    "rendering": "PASS"
}}"""
        
        llm = LLM()
        response = await llm.aask(prompt)
        response = clean_llm_response(response, expect_json=True)
        test_results = json.loads(response)
        logger.info(f"Game test results: {test_results}")
        return test_results

class GameDesigner(Role):
    """游戏设计师角色"""
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([GameDesign])
        self._watch([CodeImplementation])
        
    async def run(self, message: Message) -> Message:
        """运行游戏设计流程"""
        design_action = GameDesign()
        design_doc = await design_action.run()
        return Message(content=json.dumps(design_doc))

class SoftwareEngineer(Role):
    """软件工程师角色"""
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([CodeImplementation])
        self._watch([GameDesign, GameTester])
        
    async def run(self, message: Message) -> Message:
        """运行代码实现流程"""
        code_action = CodeImplementation()
        code = await code_action.run(json.loads(message.content))
        return Message(content=code)

class QualityEngineer(Role):
    """质量工程师角色"""
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([GameTester])
        self._watch([CodeImplementation])
        
    async def run(self, message: Message) -> Message:
        """运行测试流程"""
        test_action = GameTester()
        test_result = await test_action.run(message.content)
        return Message(content=json.dumps(test_result))

async def startup():
    """启动游戏开发项目"""
    designer = GameDesigner()
    engineer = SoftwareEngineer()
    tester = QualityEngineer()
    
    # 启动设计流程
    design_msg = Message(content="Design a CLI snake game")
    design_result = await designer.run(design_msg)
    design_doc = json.loads(design_result.content)
    
    if not design_doc:
        logger.error("Game design failed")
        return
        
    # 触发代码实现
    code_msg = Message(content=json.dumps(design_doc))
    code_result = await engineer.run(code_msg)
    final_code = code_result.content
    
    if not final_code:
        logger.error("Code implementation failed")
        return
        
    # 触发测试
    test_msg = Message(content=final_code)
    test_result = await tester.run(test_msg)
    test_data = json.loads(test_result.content)
    
    if not test_data:
        logger.error("Game testing failed")
        return
    
    if not all(result == "PASS" for result in test_data.values()):
        logger.error("Some tests failed")
        return
        
    print("\n=== Final Game Code ===")
    print(final_code)
    
    # 保存到文件
    with open("snake_game.py", "w") as f:
        f.write(final_code)
    logger.info("Game code saved to snake_game.py")

if __name__ == "__main__":
    import asyncio
    asyncio.run(startup())